# 网络类
from modules.activation import get_activation
from modules.actor_critic import *
from modules.depth_backbone import *
from modules.estimator import *
import torch

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

class model(nn.Module):
    def __init__(self, n_proprio, n_scan, n_actions, n_priv_latent, n_priv_explicit, n_hist):
        super().__init__()
        self.depth_backbone = DepthOnlyFCBackbone58x87(32)
        self.depth_model = RecurrentDepthBackbone(self.depth_backbone, n_proprio)
        self.actor_model = Actor(n_proprio, n_scan, n_actions, [128, 64, 32],
                    [512, 256, 128], [64, 20], n_priv_latent,
                    n_priv_explicit, n_hist, get_activation("elu"))
        self.estimator_model = Estimator(n_proprio, n_priv_explicit, [128, 64], "elu")
        load_path = "/home/mi/MyRL/models/model_17000.pt"
        as_state_dict = torch.load(load_path)

        self.depth_model.to(device)
        self.actor_model.to(device)
        self.estimator_model.to(device)
        self.depth_model.load_state_dict(as_state_dict['depth_encoder_state_dict'])
        self.actor_model.load_state_dict((as_state_dict['depth_actor_state_dict']))
        self.estimator_model.load_state_dict(as_state_dict['estimator_state_dict'])